﻿using System;
using System.Collections.Generic;
using System.Linq;
using Algoritmia.DataStructures.Queues;
using Algoritmia.DataStructures.Trees;
using Algoritmia.Utils;

namespace Algoritmia.Problems.Traversals
{
    public class InorderTreeTraverser<TNode,TResult> : ITreeTraverser<TNode,TResult>
    {
        private readonly Func<ILifo<IRootedTree<TNode>>> _createLifo;

        public InorderTreeTraverser(Func<ILifo<IRootedTree<TNode>>> createLifo = null)
        {
            _createLifo = createLifo ?? (() => AlgoritmiaFactory.Lifos<IRootedTree<TNode>>.Default);
        }

        #region ITreeTraverser<TNode,TResult> Members

        public IEnumerable<TResult> Traverse(IRootedTree<TNode> tree, Func<IRootedTree<TNode>, TResult> visitor = null)
        {
            var stack = _createLifo();
            var centerStack = _createLifo();
            visitor = visitor ?? (subtree => (TResult)(object) subtree.Root);
            stack.Push(tree);
            
            while ( stack.Count > 0)
            {
                var currentTree = stack.Pop();
                if ( currentTree == null)
                {
                    yield return visitor(centerStack.Pop());
                }
                else
                {
                    var subtrees = currentTree.SubTrees().ToList();
                    if ( subtrees.Count == 2)
                    {
                        stack.Push(subtrees[1]);
                    }
                    stack.Push(null);
                    centerStack.Push(currentTree);
                    if ( subtrees.Count == 2)
                    {
                        stack.Push(subtrees[0]);
                    }
                }
            }
        }

        #endregion
    }
}
